#ifndef _FLEXFLOW_SUBSTITUTIONS_SERIALIZATION_H
#define _FLEXFLOW_SUBSTITUTIONS_SERIALIZATION_H

#include "substitutions.h"

NLOHMANN_JSON_SERIALIZE_ENUM(OperatorType,
                             {{OP_INVALID, nullptr},
                              {OP_NOOP, "OP_NOOP"},
                              {OP_CONV2D, "OP_CONV2D"},
                              {OP_DROPOUT, "OP_DROPOUT"},
                              {OP_LINEAR, "OP_LINEAR"},
                              {OP_BATCHMATMUL, "OP_BATCHMATMUL"},
                              {OP_POOL2D, "OP_POOL2D_MAX"},
                              {OP_SCALAR_MULTIPLY, "OP_SCALAR_MULTIPLY"},
                              {OP_SCALAR_ADD, "OP_SCALAR_ADD"},
                              {OP_SCALAR_FLOOR_DIV, "OP_SCALAR_FLOOR_DIV"},
                              {OP_SCALAR_TRUE_DIV, "OP_SCALAR_TRUE_DIV"},
                              {OP_SCALAR_SUB, "OP_SCALAR_SUB"},
                              {OP_RELU, "OP_RELU"},
                              {OP_IDENTITY, "OP_IDENTITY"},
                              {OP_SIGMOID, "OP_SIGMOID"},
                              {OP_TANH, "OP_TANH"},
                              {OP_ELU, "OP_ELU"},
                              {OP_FLAT, "OP_FLAT"},
                              {OP_SOFTMAX, "OP_SOFTMAX"},
                              {OP_BATCHNORM, "OP_BATCHNORM"},
                              {OP_CONCAT, "OP_CONCAT"},
                              {OP_SPLIT, "OP_SPLIT"},
                              {OP_EMBEDDING, "OP_EMBEDDING"},
                              {OP_CACHE, "OP_CACHE"},
                              {OP_RESHAPE, "OP_RESHAPE"},
                              {OP_REVERSE, "OP_REVERSE"},
                              {OP_TRANSPOSE, "OP_TRANSPOSE"},
                              {OP_EW_ADD, "OP_EW_ADD"},
                              {OP_EW_MUL, "OP_EW_MUL"},
                              {OP_MATMUL, "OP_MATMUL"},
                              {OP_MUL, "OP_MUL"},
                              {OP_ENLARGE, "OP_ENLARGE"},
                              {OP_MERGE_GCONV, "OP_MERGE_GCONV"},
                              {OP_CONSTANT_IMM, "OP_CONSTANT_IMM"},
                              {OP_CONSTANT_ICONV, "OP_CONSTANT_ICONV"},
                              {OP_CONSTANT_ONE, "OP_CONSTANT_ONE"},
                              {OP_CONSTANT_POOL, "OP_CONSTANT_POOL"},
                              {OP_SQUEEZE, "OP_SQUEEZE"},
                              {OP_UNSQUEEZE, "OP_UNSQUEEZE"},
                              {OP_EW_SUB, "OP_EW_SUB"},
                              {OP_EW_DIV, "OP_EW_DIV"},
                              {OP_EW_EQUAL, "OP_EW_EQUAL"},
                              {OP_EW_GREATER, "OP_EW_GREATER"},
                              {OP_EW_LESS, "OP_EW_LESS"},
                              {OP_EW_MAX, "OP_EW_MAX"},
                              {OP_EW_MIN, "OP_EW_MIN"},
                              {OP_REDUCE_ARGMAX, "OP_REDUCE_ARGMAX"},
                              {OP_REDUCE_ARGMIN, "OP_REDUCE_ARGMIN"},
                              {OP_REDUCE_MAX, "OP_REDUCE_MAX"},
                              {OP_REDUCE_MEAN, "OP_REDUCE_MEAN"},
                              {OP_REDUCE_MIN, "OP_REDUCE_MIN"},
                              {OP_REDUCE_PROD, "OP_REDUCE_PROD"},
                              {OP_REDUCE_SUM, "OP_REDUCE_SUM"},
                              {OP_PAD, "OP_PAD"},
                              {OP_SHAPE, "OP_SHAPE"},
                              {OP_SIZE, "OP_SIZE"},
                              {OP_TOPK, "OP_TOPK"},
                              {OP_WHERE, "OP_WHERE"},
                              {OP_CEIL, "OP_CEIL"},
                              {OP_CAST, "OP_CAST"},
                              {OP_EXP, "OP_EXP"},
                              {OP_ROUND, "OP_ROUND"},
                              {OP_LOG, "OP_LOG"},
                              {OP_LOGICAL_NOT, "OP_LOGICAL_NOT"},
                              {OP_SQRT, "OP_SQRT"},
                              {OP_SIN, "OP_SIN"},
                              {OP_COS, "OP_COS"},
                              {OP_LEAKYRELU, "OP_LEAKYRELU"},
                              {OP_SLICE, "OP_SLICE"},
                              {OP_RESIZE, "OP_RESIZE"},
                              {OP_PRELU, "OP_PRELU"},
                              {OP_GELU, "OP_GELU"},
                              {OP_MULTIHEAD_ATTENTION,
                               "OP_MULTIHEAD_ATTENTION"},
                              {OP_FUSED, "OP_FUSED"},
                              {OP_RSQRT, "OP_RSQRT"},
                              {OP_POW, "OP_POW"},
                              {OP_MEAN, "OP_MEAN"},
                              {OP_LAYERNORM, "OP_LAYERNORM"},
                              {OP_REPARTITION, "OP_PARTITION"},
                              {OP_COMBINE, "OP_COMBINE"},
                              {OP_REPLICATE, "OP_REPLICATE"},
                              {OP_REDUCTION, "OP_REDUCE"},
                              {OP_PIPELINE, "OP_PIPELINE"},
                              {OP_FUSED_PARALLEL, "OP_FUSED_PARALLEL"}})

namespace FlexFlow {
namespace substitutions {

NLOHMANN_JSON_SERIALIZE_ENUM(
    ParameterAttribute,
    {{ParameterAttribute::INVALID, nullptr},
     {ParameterAttribute::OP_TYPE, "PM_OP_TYPE"},
     {ParameterAttribute::NUM_INPUTS, "PM_NUM_INPUTS"},
     {ParameterAttribute::NUM_OUTPUTS, "PM_NUM_OUTPUTS"},
     {ParameterAttribute::GROUP, "PM_GROUP"},
     {ParameterAttribute::KERNEL_H, "PM_KERNEL_H"},
     {ParameterAttribute::KERNEL_W, "PM_KERNEL_W"},
     {ParameterAttribute::STRIDE_H, "PM_STRIDE_H"},
     {ParameterAttribute::STRIDE_W, "PM_STRIDE_W"},
     {ParameterAttribute::PADDING_H, "PM_PADDING_H"},
     {ParameterAttribute::PADDING_W, "PM_PADDING_W"},
     {ParameterAttribute::ACTIVATION, "PM_ACTIVATION"},
     {ParameterAttribute::NUMDIM, "PM_NUMDIM"},
     {ParameterAttribute::AXIS, "PM_AXIS"},
     {ParameterAttribute::PERM, "PM_PERM"},
     {ParameterAttribute::OUTSHUFFLE, "PM_OUTSHUFFLE"},
     {ParameterAttribute::MERGE_GCONV_COUNT, "PM_MERGE_GCONV_COUNT"},
     {ParameterAttribute::AXES, "PM_AXES"},
     {ParameterAttribute::KEEP_DIMS, "PM_KEEP_DIMS"},
     {ParameterAttribute::EPSILON, "PM_EPSILON"},
     {ParameterAttribute::REPARTITION_DIM, "PM_REPARTITION_DIM"},
     {ParameterAttribute::REPARTITION_DEGREE, "PM_REPARTITION_DEGREE"},
     {ParameterAttribute::REPLICATE_DIM, "PM_REPLICATE_DIM"},
     {ParameterAttribute::REPLICATE_DEGREE, "PM_REPLICATE_DEGREE"},
     {ParameterAttribute::COMBINE_DIM, "PM_COMBINE_DIM"},
     {ParameterAttribute::COMBINE_DEGREE, "PM_COMBINE_DEGREE"},
     {ParameterAttribute::REDUCTION_DIM, "PM_REDUCTION_DIM"},
     {ParameterAttribute::REDUCTION_DEGREE, "PM_REDUCTION_DEGREE"},
     {ParameterAttribute::SOFTMAX_DIM, "PM_SOFTMAX_DIM"},
     {ParameterAttribute::NUM_HEADS, "PM_NUM_HEADS"},
     {ParameterAttribute::PARALLEL_DIM, "PM_PARALLEL_DIM"},
     {ParameterAttribute::PARALLEL_DEGREE, "PM_PARALLEL_DEGREE"},
     {ParameterAttribute::PAD, "PM_PAD"}})
void from_json(nlohmann::json const &j, OperatorAttributeConstraint &p);
void from_json(nlohmann::json const &j, Tensor &t);
void from_json(nlohmann::json const &j, OperatorConstraint &t);
void from_json(nlohmann::json const &j, MapOutput &t);
void from_json(nlohmann::json const &j, Rule &t);
void from_json(nlohmann::json const &j, RuleCollection &c);

RuleCollection load_rule_collection(std::istream &s);
RuleCollection load_rule_collection_from_path(std::string const &path);

} // namespace substitutions
} // namespace FlexFlow

#endif
